package cn.iocoder.yudao.server.controller;

import cn.iocoder.yudao.server.utils.AIRecognizeUtil;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.http.ResponseEntity;
import org.springframework.web.bind.annotation.*;
import org.springframework.web.multipart.MultipartFile;

import java.io.IOException;
import java.io.InputStream;

@RestController
@RequestMapping("/api/ai-recognize")
public class AIRecognizeController {
    private static final Logger log = LoggerFactory.getLogger(AIRecognizeController.class);

    @PostMapping("/certificate")
    public ResponseEntity<String> recognizeCertificate(@RequestPart("file") MultipartFile file) {
        try {
            // 1. 检查文件是否为空
            if (file.isEmpty()) {
                log.warn("上传文件为空");
                return ResponseEntity.badRequest().body("{\"error\": \"文件不能为空\"}");
            }

            // 2. 根据文件类型处理
            byte[] imageBytes;
            String contentType = file.getContentType();

            if (contentType != null && contentType.equals("application/pdf")) {
                // PDF文件：转换为图片
                try (InputStream pdfStream = file.getInputStream()) {
                    imageBytes = AIRecognizeUtil.convertPdfToImageBytes(pdfStream);
                }
                // 3. 调用OCR识别
                String ocrResult = AIRecognizeUtil.recognizeImage(imageBytes);
                log.info("OCR识别成功，结果长度: {}", ocrResult.length());
                return buildResponse(ocrResult);
            } else if (contentType != null && contentType.startsWith("image/")) {
                // 图片文件：直接获取字节
                imageBytes = file.getBytes();
                // 3. 调用OCR识别
                String ocrResult = AIRecognizeUtil.recognizeImage(imageBytes);
                log.info("OCR识别成功，结果长度: {}", ocrResult.length());
                return buildResponse(ocrResult);
            } else if (contentType != null && (contentType.equals("application/msword") || file.getOriginalFilename().endsWith(".doc"))) {
                // doc 文件
                String text;
                try (InputStream in = file.getInputStream()) {
                    text = AIRecognizeUtil.extractTextFromDoc(in);
                }
                String ocrResult = AIRecognizeUtil.recognizeText(text);
                log.info("文本识别成功，结果长度: {}", ocrResult.length());
                return buildResponse(ocrResult);
            } else if (contentType != null && (contentType.equals("application/vnd.openxmlformats-officedocument.wordprocessingml.document") || file.getOriginalFilename().endsWith(".docx"))) {
                // docx 文件
                String text;
                try (InputStream in = file.getInputStream()) {
                    text = AIRecognizeUtil.extractTextFromDocx(in);
                }
                String ocrResult = AIRecognizeUtil.recognizeText(text);
                log.info("文本识别成功，结果长度: {}", ocrResult.length());
                return buildResponse(ocrResult);
            } else if (contentType != null && contentType.equals("text/plain")) {
                // txt 文件
                String text;
                try (InputStream in = file.getInputStream()) {
                    text = AIRecognizeUtil.extractTextFromTxt(in);
                }
                String ocrResult = AIRecognizeUtil.recognizeText(text);
                log.info("文本识别成功，结果长度: {}", ocrResult.length());
                return buildResponse(ocrResult);
            } else {
                log.warn("不支持的文件类型: {}", contentType);
                return ResponseEntity.badRequest().body("{\"error\": \"只支持图片、PDF、Word、TXT文件\"}");
            }
        } catch (IOException e) {
            log.error("文件处理失败", e);
            return ResponseEntity.internalServerError().body("{\"error\": \"文件处理失败: " + e.getMessage() + "\"}");
        } catch (Exception e) {
            log.error("OCR识别失败", e);
            return ResponseEntity.internalServerError().body("{\"error\": \"识别服务异常: " + e.getMessage() + "\"}");
        }
    }

    // 新增：统一的结果解析方法
    private static ResponseEntity<String> buildResponse(String ocrResult) {
        try {
            org.json.JSONObject resultJson = new org.json.JSONObject(ocrResult);
            org.json.JSONObject responseJson = new org.json.JSONObject();
            org.json.JSONObject dataJson;

            if (resultJson.has("output")) {
                // 图片接口结构
                org.json.JSONObject output = resultJson.getJSONObject("output");
                org.json.JSONArray choicesArray = output.getJSONArray("choices");
                org.json.JSONObject choice = choicesArray.getJSONObject(0);
                org.json.JSONObject message = choice.getJSONObject("message");
                org.json.JSONArray contentArray = message.getJSONArray("content");
                org.json.JSONObject messageContent = contentArray.getJSONObject(0);
                String jsonText = messageContent.getString("text").replace("```json", "").replace("```", "");
                dataJson = new org.json.JSONObject(jsonText);
            } else if (resultJson.has("choices")) {
                // 文本接口结构
                org.json.JSONArray choicesArray = resultJson.getJSONArray("choices");
                org.json.JSONObject choice = choicesArray.getJSONObject(0);
                org.json.JSONObject message = choice.getJSONObject("message");
                String content = message.getString("content");
                // 只提取大括号内容
                int start = content.indexOf("{");
                int end = content.lastIndexOf("}") + 1;
                String jsonText = (start >= 0 && end > start) ? content.substring(start, end) : content;
                dataJson = new org.json.JSONObject(jsonText);
            } else {
                throw new Exception("未知的AI返回结构: " + ocrResult);
            }

            responseJson.put("code", 200);
            responseJson.put("data", dataJson);
            responseJson.put("message", "识别成功");
            return ResponseEntity.ok(responseJson.toString());
        } catch (Exception e) {
            log.error("结果解析失败", e);
            return ResponseEntity.internalServerError().body("{\"error\": \"结果解析失败: " + e.getMessage() + "\"}");
        }
    }
}